// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <map>
#include <string>
#include <pollux/functions/macros.h>
#include <pollux/common/strings/string_core.h>
#include <pollux/type/decimal_util.h>

namespace kumo::pollux::functions {
    enum class Unit : int128_t {
        BYTE = 1,
        KILOBYTE = int128_t(1) << 10,
        MEGABYTE = int128_t(1) << 20,
        GIGABYTE = int128_t(1) << 30,
        TERABYTE = int128_t(1) << 40,
        PETABYTE = int128_t(1) << 50,
        EXABYTE = int128_t(1) << 60,
        ZETTABYTE = int128_t(1) << 70,
        YOTTABYTE = int128_t(1) << 80
    };

    inline Unit parseUnit(const std::string &dataSize, const size_t &valueLength) {
        static const std::map<std::string, Unit> unitMap = {
            {"B", Unit::BYTE},
            {"kB", Unit::KILOBYTE},
            {"MB", Unit::MEGABYTE},
            {"GB", Unit::GIGABYTE},
            {"TB", Unit::TERABYTE},
            {"PB", Unit::PETABYTE},
            {"EB", Unit::EXABYTE},
            {"ZB", Unit::ZETTABYTE},
            {"YB", Unit::YOTTABYTE}
        };
        try {
            std::string unitString = dataSize.substr(valueLength);
            auto it = unitMap.find(unitString);
            POLLUX_USER_CHECK(it != unitMap.end(), "Invalid data size: '{}'", dataSize);
            return it->second;
        } catch (const std::exception &) {
            POLLUX_USER_FAIL("Invalid data size: '{}'", dataSize);
        }
    }

    inline double parseValue(
        const std::string &dataSize,
        const size_t &valueLength) {
        try {
            std::string value = dataSize.substr(0, valueLength);
            return std::stod(value);
        } catch (const std::exception &) {
            POLLUX_USER_FAIL("Invalid data size: '{}'", dataSize);
        }
    }

    inline int128_t getDecimal(const std::string &dataSize) {
        size_t valueLength = 0;
        while (valueLength < dataSize.length() &&
               (isdigit(dataSize[valueLength]) || dataSize[valueLength] == '.')) {
            valueLength++;
        }
        POLLUX_USER_CHECK_GT(valueLength, 0, "Invalid data size: '{}'", dataSize);
        double value = parseValue(dataSize, valueLength);
        Unit unit = parseUnit(dataSize, valueLength);

        int128_t factor = static_cast<int128_t>(unit);
        int128_t scaledValue = static_cast<int128_t>(value * factor);

        // Ensure the result is within a valid range.
        try {
            DecimalUtil::valueInRange(scaledValue);
            std::vector<char> encodedValue(
                DecimalUtil::getByteArrayLength(scaledValue));
            DecimalUtil::toByteArray(scaledValue, encodedValue.data());
        } catch (const std::exception &) {
            POLLUX_USER_FAIL("Value out of range: '{}' ('{}B')", dataSize, scaledValue);
        }
        return scaledValue;
    }

    template<typename TExec>
    struct ParsePrestoDataSizeFunction {
        POLLUX_DEFINE_FUNCTION_TYPES(TExec);
        static constexpr bool is_default_ascii_behavior = true;
        MELON_ALWAYS_INLINE void callAscii(
            out_type<LongDecimal<P1, S1> > &result,
            const arg_type<Varchar> &input) {
            result = getDecimal(input);
        }

        // Fail Non-ASCII characters in input.
        MELON_ALWAYS_INLINE void call(
            out_type<LongDecimal<P1, S1> > &result,
            const arg_type<Varchar> &input) {
            POLLUX_USER_FAIL("Invalid data size: '{}'", input);
        }
    };
} // namespace kumo::pollux::functions
